import os
import torch
from torch import nn
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from model.PResnet_model import PResnet
from params import params_dict


def load_model(params_dict):
    model = PResnet(params_dict['vocab_size'], params_dict['embedding_dim'], params_dict['n_head'], params_dict['num_layers'],
                 params_dict['v_vector_size'], params_dict['input_dim'], params_dict['output_dim'], params_dict['is_train'])
    if params_dict['pre_train'] == True:
        model.load_state_dict(torch.load(params_dict['pre_train_model_path']))
    return model


def save_model(model, save_file_path, save_model_path):
    """
    保存模型权重
    :param save_path: 模型保存路径
    :param model: 模型权重
    :return:
    """
    if not os.path.exists(save_file_path):
        os.mkdir(save_file_path)
    torch.save(model.state_dict(), save_model_path)
    return None


def set_optimizer(model, lr):
    """
    设置优化器
    :param model: 模型
    :param lr: 学习率
    :return:
    """
    optimizer = torch.optim.Adam(model.parameters(), lr)
    # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer,
    #                                                        T_max=params_dict['epoch'])
    return optimizer


def set_loss():
    """
    设置交叉熵损失
    :return:
    """
    return nn.CrossEntropyLoss()


# if __name__ == '__main__':
#     model = load_model(params_dict)
#     save_model(model, "../run/PNAT2024-05-31 10_36_40", os.path.join("../run/PNAT2024-05-31 10_36_40", "best.pt"))
